#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include "mpi.h"

#define FILTERDIM 5
#define FILTERSIZE 25

#define MASTER		0
#define WORKER_R	1
#define WORKER_G	2
#define WORKER_B	3

#define RED 	'r'
#define GREEN 	'g'
#define BLUE 	'b'

struct RGBstructure {
    int r;
    int g;
    int b;
};

typedef struct RGBstructure RGB;

struct PPMstructure {
  char type[2];
  int width;
  int height;
  int maxVal;
  RGB **data;
};

typedef struct PPMstructure PPMImage;

void allocData(PPMImage *img, char* type, int width, int height, int maxVal) {
    int i;

    strcpy(img->type, type);
    img->width = width;
    img->height = height;
    img->maxVal = maxVal;

    img->data = (RGB **)malloc(height * sizeof(RGB*));
    for (i = 0; i < height; i++) {
        img->data[i] = (RGB *)malloc(width * sizeof(RGB));
    }
}

void allocDataChannel(PPMImage *img, int width, int height) {
    int i;

    img->width = width;
    img->height = height;

    img->data = (RGB **)malloc(height * sizeof(RGB*));

    for (i = 0; i < height; i++) {
        img->data[i] = (RGB *)malloc(width * sizeof(RGB));
    }
}

void freeData(PPMImage *img) {
    int i;

    for (i = 0; i < img->height; i++) {
        free(img->data[i]);
    }

    free(img->data);
}

/* Cria estrutura PPMImage */
void readImage(PPMImage *imgIn, PPMImage *imgOut, char *in) {
    FILE *input;
    int i, j, ch;
    char type[2];
    int maxVal, width, height;

    input = fopen(in, "r");

    if (input == NULL)
        printf("error\n");
    else {
        /* Le header */
        // Le tipo
        fscanf(input, "%s", type);
        // Pula fim da linha
        while (getc(input) != '\n');

        // Pula comentario
        while (getc(input) == '#') {
            while (getc(input) != '\n');
        }

        // Volta um caracter
        fseek(input, -1, SEEK_CUR);

        // Le tamanho
        fscanf(input, "%d", &width);
        fscanf(input, "%d", &height);
        fscanf(input, "%d", &maxVal);

        allocData(imgIn, type, width, height, maxVal);
        allocData(imgOut, type, width, height, maxVal);
        
        // Le dados
        for (i = 0; i < height; i++) {
            for (j = 0; j < width; j++) {
                fscanf(input, "%d", &ch);
                imgIn->data[i][j].r = ch;
                fscanf(input, "%d", &ch);
                imgIn->data[i][j].g = ch;
                fscanf(input, "%d", &ch);
                imgIn->data[i][j].b = ch;
            }
        }
    }

    fclose(input);
}

/* Cria arquivo com o resultado */
void saveImage(PPMImage *img, char *out) {
    FILE *output;
    int i, j;

    // Escreve header
    output = fopen(out, "w");
    fprintf(output, "%s\n", img->type);
    fprintf(output, "%d %d\n", img->width, img->height);
    fprintf(output, "%d\n", img->maxVal);

    // Escreve dados
    for (i = 0; i < img->height; i++) {
        for (j = 0; j < img->width; j++) {
            fprintf(output, "%d %d %d\t", 
                img->data[i][j].r,
                img->data[i][j].g,
                img->data[i][j].b);
        }
        fprintf(output, "\n");
    }

    fclose(output);
}

void MPIfilter(PPMImage *imgIn, PPMImage *imgOut, int processId) {
    MPI_Status status;

    // Valor para controle do indice do filtro
    int a = (FILTERDIM - 1) / 2; 

    if (processId == MASTER){
		int x,y;
		// Copia bordas da imagem original para evitar tratamento de bordas
		for (x = 0; x < imgIn->height; x++) {
			for (y = 0; y < a; y++)
				imgOut->data[x][y] = imgIn->data[x][y];
		}
		for (x = 0; x < imgIn->height; x++) {
			for (y = imgIn->width - a; y < imgIn->width; y++)
				imgOut->data[x][y] = imgIn->data[x][y];
		}
		for (x = 0; x < a; x++) {
			for (y = a; y < imgIn->width - a; y++)
				imgOut->data[x][y] = imgIn->data[x][y];
		}
		for (x = imgIn->height - a; x < imgIn->height; x++) {
			for (y = a; y < imgIn->width - a; y++)
				imgOut->data[x][y] = imgIn->data[x][y];
		}

		//Recebe pixels filtrados
		for (x = a; x < imgIn->height - a; x++) {
			for (y = a; y < imgIn->width - a; y++){
				MPI_Recv(&(imgOut->data[x][y].r), 1, MPI_INT, WORKER_R, RED + 2, MPI_COMM_WORLD, &status);
				MPI_Recv(&(imgOut->data[x][y].g), 1, MPI_INT, WORKER_G, GREEN + 2, MPI_COMM_WORLD, &status);
				MPI_Recv(&(imgOut->data[x][y].b), 1, MPI_INT, WORKER_B, BLUE + 2, MPI_COMM_WORLD, &status);
			}
		}

     	return;
    }

    int x, y, s, t;
    float filter[FILTERDIM][FILTERDIM];
    // Cria filtro de media
    for (x = 0; x < FILTERDIM; x++) {
        for(y = 0; y < FILTERDIM; y++) {
            filter[x][y] = 1.0 / (float)(FILTERSIZE);
        }
    }

	if (processId == WORKER_R) {

		for (x = a; x < imgIn->height - a; x++) {
			for(y = a; y < imgIn->width - a; y++) {
				float newvalueR = 0;

				for (s = -a; s <= a; s++) {
					for (t = -a ; t <= a; t++) {
						float valueR = imgIn->data[x + s][y + t].r;
        				newvalueR += filter[s + a][t + a] * valueR;
        			}   
				}
				// Envia o pixel ao mestre
				int newR = (int)newvalueR;
				MPI_Send(&newR, 1, MPI_INT, MASTER, RED + 2, MPI_COMM_WORLD);
			}
		}
	}
	if (processId == WORKER_G) {

		for (x = a; x < imgIn->height - a; x++) {
			for(y = a; y < imgIn->width - a; y++) {
				float newvalueG = 0;

				for (s = -a; s <= a; s++) {
					for (t = -a ; t <= a; t++) {
						float valueG = imgIn->data[x + s][y + t].g;
        				newvalueG += filter[s + a][t + a] * valueG;
        			}   
				}

				int newG = (int)newvalueG;
				MPI_Send(&newG, 1, MPI_INT, MASTER, GREEN + 2, MPI_COMM_WORLD);
			}
		}
	}
	if (processId == WORKER_B) {
		for (x = a; x < imgIn->height - a; x++) {
			for(y = a; y < imgIn->width - a; y++) {
				float newvalueB = 0;

				for (s = -a; s <= a; s++) {
					for (t = -a ; t <= a; t++) {
						float valueB = imgIn->data[x + s][y + t].b;
        				newvalueB += filter[s + a][t + a] * valueB;
        			}   
				}

				int newB = (int)newvalueB;
				MPI_Send(&newB, 1, MPI_INT, MASTER, BLUE + 2, MPI_COMM_WORLD);
			}
		}
	}
} 

int main(int argc, char *argv[]) { 
    PPMImage imgIn, imgOut;
    // Guarda o ID do processo corrente (rank)
    int processId;
    int i,j;	
    char in[20], out[20], tempo[20];

	// Tempo total
    struct timeval startTimeMPITotal, endTimeMPITotal;
    gettimeofday(&startTimeMPITotal, NULL);

    strcpy(in, argv[1]);
    strcpy(out, argv[2]);
    strcpy(tempo, argv[3]);			
    // Inicializa MPI
    MPI_Status status;
    MPI_Init(&argc, &argv);

    // Le o rank
    MPI_Comm_rank(MPI_COMM_WORLD, &processId);
    

	if (processId == MASTER){
		readImage(&imgIn, &imgOut, in);
		int *red = (int *)malloc(imgIn.height * imgIn.width * sizeof(int));
		int *green = (int *)malloc(imgIn.height * imgIn.width * sizeof(int));
		int *blue = (int *)malloc(imgIn.height * imgIn.width * sizeof(int));

		for(i = 0; i < imgIn.height; i++) {
			for(j = 0; j < imgIn.width; j++) {
				red[(i * imgIn.width) + j] = imgIn.data[i][j].r;
				green[(i * imgIn.width) + j] = imgIn.data[i][j].g;
				blue[(i * imgIn.width) + j] = imgIn.data[i][j].b;
			}	
		}

		// Manda o header correspondente por canal
		MPI_Send(&(imgIn.width), 1, MPI_INT, WORKER_R, RED, MPI_COMM_WORLD);
		MPI_Send(&(imgIn.height), 1, MPI_INT, WORKER_R, RED + 1, MPI_COMM_WORLD);
		MPI_Send(&(red[0]), imgIn.width * imgIn.height, MPI_INT, WORKER_R, RED + 2, MPI_COMM_WORLD);

		MPI_Send(&(imgIn.width), 1, MPI_INT, WORKER_G, GREEN, MPI_COMM_WORLD);
		MPI_Send(&(imgIn.height), 1, MPI_INT, WORKER_G, GREEN + 1, MPI_COMM_WORLD);
		MPI_Send(&(green[0]), imgIn.width * imgIn.height, MPI_INT, WORKER_G, GREEN + 2, MPI_COMM_WORLD);

		MPI_Send(&(imgIn.width), 1, MPI_INT, WORKER_B, BLUE, MPI_COMM_WORLD);
		MPI_Send(&(imgIn.height), 1, MPI_INT, WORKER_B, BLUE + 1, MPI_COMM_WORLD);
		MPI_Send(&(blue[0]), imgIn.width * imgIn.height, MPI_INT, WORKER_B, BLUE + 2, MPI_COMM_WORLD);

	} else if (processId == WORKER_R){
		int width,height;
		// Recebe os dados
		MPI_Recv(&(width), 1, MPI_INT, MASTER, RED, MPI_COMM_WORLD, &status);
		MPI_Recv(&(height), 1, MPI_INT, MASTER, RED + 1, MPI_COMM_WORLD, &status);

		int *red = (int *)malloc(height * width * sizeof(int));

		MPI_Recv(&(red[0]), width * height, MPI_INT, MASTER, RED + 2, MPI_COMM_WORLD, &status);

		allocDataChannel(&imgIn, width, height);		
		allocDataChannel(&imgOut, width, height);

		for( i = 0; i < height; i++) {
			for( j = 0; j < width; j++) {
				imgIn.data[i][j].r = red[(i * imgIn.width) + j];
			}	
		}
	} else if (processId == WORKER_G){
		int width,height;
		
		MPI_Recv(&(width), 1, MPI_INT, MASTER, GREEN, MPI_COMM_WORLD, &status);
		MPI_Recv(&(height), 1, MPI_INT, MASTER, GREEN + 1, MPI_COMM_WORLD, &status);

		int *green = (int *)malloc(height * width * sizeof(int));

		MPI_Recv(&(green[0]), width * height, MPI_INT, MASTER, GREEN + 2, MPI_COMM_WORLD, &status);

		allocDataChannel(&imgIn, width, height);
		allocDataChannel(&imgOut, width, height);

		for( i = 0; i < height; i++) {
			for( j = 0; j < width; j++) {
				imgIn.data[i][j].g = green[(i * imgIn.width) + j];
			}	
		}
	} else if (processId == WORKER_B){
		int width,height;

		MPI_Recv(&(width), 1, MPI_INT, MASTER, BLUE, MPI_COMM_WORLD, &status);
		MPI_Recv(&(height), 1, MPI_INT, MASTER, BLUE + 1, MPI_COMM_WORLD, &status);
		
		int *blue = (int *)malloc(height * width * sizeof(int));
		
		MPI_Recv(&(blue[0]), width * height, MPI_INT, MASTER, BLUE + 2, MPI_COMM_WORLD, &status);
		
		allocDataChannel(&imgIn, width, height);		
		allocDataChannel(&imgOut, width, height);
		
		for( i = 0; i < height; i++) {
			for( j = 0; j < width; j++) {
				imgIn.data[i][j].b = blue[(i * imgIn.width) + j];
			}	
		}
	} else {
		MPI_Finalize();
		return 0;
	}

	// Tempo do smoothing
    struct timeval startTimeMPI, endTimeMPI;
    gettimeofday(&startTimeMPI, NULL);
   
    MPIfilter(&imgIn, &imgOut, processId);

    gettimeofday(&endTimeMPI, NULL);

    if (processId == MASTER) {
        saveImage(&imgOut, out);
        freeData(&imgIn);
    }
    
    freeData(&imgOut);
    
    MPI_Finalize();

    gettimeofday(&endTimeMPITotal, NULL); 

    // Calcula o tempo
    double result = endTimeMPI.tv_sec - startTimeMPI.tv_sec + (endTimeMPI.tv_usec - startTimeMPI.tv_usec) / 1000000.0;
    //double resultTotal = endTimeMPITotal.tv_sec - startTimeMPITotal.tv_sec + (endTimeMPITotal.tv_usec - startTimeMPITotal.tv_usec) / 1000000.0;

    printf("MPI: %lfs\n", result);

    FILE *t;
    t = fopen(tempo,"w");
    fprintf(t,"MPI: %lfs\n", result);
    fclose(t);
    
    return 0;        
}